import argparse  # 用于处理命令行参数
import gym  # 引入OpenAI Gym库，提供强化学习环境
import numpy as np  # 用于数值计算
from itertools import count  # 用于创建无限计数器

import torch  # 引入PyTorch库，用于深度学习
import torch.nn as nn  # 引入神经网络模块
import torch.nn.functional as F  # 引入常用的神经网络函数，如激活函数
import torch.optim as optim  # 引入优化器模块
from torch.distributions import Categorical  # 引入分类分布，用于策略采样

# 定义策略网络
class Policy(nn.Module):
    def __init__(self):
        super(Policy, self).__init__()
        self.affine1 = nn.Linear(4, 128)  # 输入层：状态维度为4，隐层维度为128
        self.affine2 = nn.Linear(128, 2)  # 输出层：动作维度为2（左右移动）

        self.saved_log_probs = []  # 保存动作对应的log概率，用于后续梯度计算
        self.rewards = []  # 保存回合奖励

    def forward(self, x):
        x = F.relu(self.affine1(x))  # 隐层使用ReLU激活函数
        action_scores = self.affine2(x)  # 输出动作得分
        return F.softmax(action_scores, dim=1)  # 使用Softmax将动作得分转换为概率分布


# 动作选择函数
def select_action(state):
    state = torch.from_numpy(state).float().unsqueeze(0)  # 将状态转换为浮点张量并增加一个批量维度
    probs = policy(state)  # 获取策略网络输出的动作概率
    m = Categorical(probs)  # 使用分类分布生成分布对象
    action = m.sample()  # 从分布中采样动作
    policy.saved_log_probs.append(m.log_prob(action))  # 保存动作的log概率
    return action.item()  # 返回采样的动作

# 完成一个回合并更新策略
def finish_episode():
    R = 0  # 初始化累计折扣奖励
    policy_loss = []  # 用于保存策略损失
    rewards = []  # 保存折扣奖励

    # 计算折扣奖励
    for r in policy.rewards[::-1]:  # 倒序遍历每一步的奖励
        R = r + args.gamma * R  # 计算累计奖励
        rewards.insert(0, R)  # 将累计奖励插入到列表开头

    rewards = torch.tensor(rewards)  # 将奖励转换为张量
    rewards = (rewards - rewards.mean()) / (rewards.std() + eps)  # 标准化奖励

    # 计算策略损失
    for log_prob, reward in zip(policy.saved_log_probs, rewards):
        policy_loss.append(-log_prob * reward)  # 损失是负的log概率乘以奖励

    optimizer.zero_grad()  # 清零梯度
    policy_loss = torch.cat(policy_loss).sum()  # 合并所有损失并求和
    policy_loss.backward()  # 反向传播计算梯度
    optimizer.step()  # 使用优化器更新网络参数

    del policy.rewards[:]  # 清空回合奖励
    del policy.saved_log_probs[:]  # 清空log概率

# 主函数
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch REINFORCE example')  # 创建命令行参数解析器
    parser.add_argument('--gamma', type=float, default=0.99, metavar='G',
                        help='discount factor (default: 0.99)')  # 折扣因子
    parser.add_argument('--seed', type=int, default=543, metavar='N',
                        help='random seed (default: 543)')  # 随机种子
    parser.add_argument('--render', action='store_true',
                        help='render the environment')  # 是否渲染环境
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='interval between training status logs (default: 10)')  # 日志间隔
    args = parser.parse_args()  # 解析命令行参数

    policy = Policy()  # 创建策略网络
    optimizer = optim.Adam(policy.parameters(), lr=1e-2)  # 定义优化器，使用Adam，学习率为0.01
    eps = np.finfo(np.float32).eps.item()  # 防止浮点数精度问题的小数值

    env = gym.make('CartPole-v1')  # 创建CartPole-v1环境
    torch.manual_seed(args.seed)  # 设置随机种子
    running_reward = 10  # 初始化运行奖励

    for i_episode in count(1):  # 无限循环，直到达到停止条件
        state, _ = env.reset()  # 重置环境，获取初始状态
        for t in range(10000):  # 限制最大时间步，防止无限循环
            action = select_action(state)  # 选择动作
            state, reward, done, _, _ = env.step(action)  # 执行动作，返回下一状态和奖励
            if args.render:  # 如果启用了渲染
                env.render()  # 渲染环境
            policy.rewards.append(reward)  # 保存奖励
            if done:  # 如果当前回合结束
                break  # 跳出循环

        running_reward = running_reward * 0.99 + t * 0.01  # 更新运行奖励
        finish_episode()  # 完成回合并更新策略

        if i_episode % args.log_interval == 0:  # 每隔一定回合打印日志
            print('Episode {}\tLast length: {:5d}\tAverage length: {:.2f}'.format(
                i_episode, t, running_reward))

        if running_reward > env.spec.reward_threshold:  # 如果运行奖励超过环境要求的阈值
            print("Solved! Running reward is now {} and "
                  "the last episode runs to {} time steps!".format(running_reward, t))
            break  # 训练结束
